from selenium import webdriver
from selenium.common.exceptions import TimeoutException
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from PIL import Image
from io import BytesIO
import time

# 屏幕截图
def get_screen_shot(browser):
    # 全屏截图，生成一个图片流
    screenshot = browser.get_screenshot_as_png()
    # 生成image对象，便于后面裁剪
    image = Image.open(BytesIO(screenshot))
    return image

# 取验证码坐标[左上角，右下角]
def get_captcha_position(browser, css):
    wait = WebDriverWait(browser, 5)
    captcha = wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, css)))
    left_x = captcha.location['x']
    left_y = captcha.location['y']

    width = captcha.size['width']
    height = captcha.size['height']

    print(width, height)

    right_x = left_x + width
    right_y = left_y + height

    left_x *= 1.24
    left_y *= 1.2
    right_x *= 1.26
    right_y *= 1.26

    return (left_x, left_y, right_x, right_y)

# 截取图保存
def crop_image(browser, css, image_name):
    # 获取全图
    full_image = get_screen_shot(browser)
    # 获取验证图片坐标
    left_x, left_y, right_x, right_y = get_captcha_position(browser, css)
    # 截图并保存
    captcha_image = full_image.crop((left_x, left_y, right_x, right_y))
    captcha_image.save(image_name)
    return captcha_image


